import os
import numpy
import matplotlib
matplotlib.use('agg')
from matplotlib import pyplot as PLT
from datetime import date,datetime,time

"""
    Bins, averages, and plots torque for all runs in a given month. Also plots all monthly
    averages after Feb 2012 together for comparison.
"""


def MonthlyAverages(database, bins, max_d, outdir = 'averagePlots/'):
    # get today's date for file title
    d = datetime.now()
    q = str(d.strftime('%Y%m%d'))
    
    # retrieve all month entries and store as dictionary keys
    months = {}
    for run in database:
        month_year = run.Datetime().strftime('%b%Y')
        if month_year not in months.keys():
            months[month_year] = []
    
    # bin and average data
    for m in months:
        print m
        x_inputs = numpy.zeros(bins)
        X = numpy.linspace(0,max_d,bins)
        Y = numpy.zeros(bins)
        Y2 = []
        for i in range(X.size):
            Y2.append([])
        
        # append data of each file to appropriate month dictionary entry
        for run in database:
            if run.Direction() != 'in': continue #Skip all retractions
            if run.Datetime().strftime('%b%Y') not in m: continue
            if run.Source() != 'weakTh-228': continue
            if run.Station() in ['S8', 'S11', 'Other']: continue
            print run
            position = run.NewPosition()
            torque = run.Torque()
            for j, x in enumerate(X):
                pos_torque = []
                x0 = x-(max_d/bins)/2.
                x1 = x+(max_d/bins)/2.
                for i, p in enumerate(position):
                    if p > x0 and p <= x1:
                        pos_torque.append(torque[i])
                if len(pos_torque)>0:
                    if max(pos_torque) == 0: continue
                    Y[j] += max(pos_torque)
                    Y2[j].append(max(pos_torque))
                    x_inputs[j]+=1.
        
        X_nonull = []
        Y_nonull = []
        x_inputs_nonull = []
        Y2_nonull = []
        
        # 'safety' measure in case of blank spaces
        for g, h in enumerate(x_inputs):
            if h == 0: continue
            X_nonull.append(X[g])
            Y_nonull.append(Y[g])
            x_inputs_nonull.append(x_inputs[g])
            Y2_nonull.append(Y2[g])
        #print X_nonull
        Y_std = numpy.zeros(len(X_nonull))
        
        # average content of each bin, calculate standard deviation
        for l, n in enumerate(x_inputs_nonull):
            if n==0: continue
            Y_nonull[l] = Y_nonull[l]/n
            Y_std[l] = numpy.std(Y2_nonull[l])
        months[m] = [X_nonull, Y_nonull]
        
        #_____plot individual monthly averages)_____#
        
        # checks in folder where plots are dropped to see if data has already been plotted
        alreadyPlotted = []
        dropFolder = os.listdir(outdir)
        for file in dropFolder:
            alreadyPlotted.append(file.split('.')[0].split('_')[-1])
        
        if run.Datetime().strftime('%m%Y') in alreadyPlotted:
            continue
        
        fig = PLT.figure(figsize=(15,8),dpi=150)
        box = [0.14, 0.14, 0.76, 0.76]
        ax1 = fig.add_axes(box)
        ax1.set_xlabel('Position (cm)')
        ax1.set_ylabel('Torque (in.oz)')
        miny = min(Y_std)
        maxy = max(Y)
        maxx = max(X)
            
        ## Add station locations
        if maxx >= 328.9:
            ax1.axvline(338.9, color = 'r', linewidth = 1, linestyle = '-')
            ax1.annotate('S2', xy = (340.9, miny * 0.8), color = 'r')
        if maxx >= 382.6:
            ax1.axvline(392.6, color = 'r', linewidth = 1, linestyle = '-')
            ax1.annotate('S5', xy = (394.6, miny * 0.8), color = 'r')
        if maxx >= 435.7:
            ax1.axvline(445.7, color = 'r', linewidth = 1, linestyle = '-')
            ax1.annotate('S8', xy = (447.7, miny * 0.8), color = 'r')
        if maxx >= 491.7:
            ax1.axvline(501.7, color = 'r', linewidth = 1, linestyle = '-')
            ax1.annotate('S11', xy = (503.7, miny * 0.8), color = 'r')
        if maxx >= 595.8:
            ax1.axvline(605.8, color = 'r', linewidth = 1, linestyle = '-')
            ax1.annotate('S17', xy = (607.8, miny * 0.8), color = 'r')
        
        ## Add bend locations
        if maxx >= 47.8:
            ax1.axvspan(47.8, 93.5, color = 'g', alpha = 0.1)
            ax1.annotate('Bellows', xy = (48, miny), color = 'g')
        if maxx >= 136.4:
            ax1.axvspan(136.4, 144.6, color = 'g', alpha = 0.1)
            ax1.annotate('First outer bend', xy = (136, miny), color = 'g')
        if maxx >= 196.5:
            ax1.axvspan(196.5, 218.0, color = 'g', alpha = 0.1)
            ax1.annotate('The "S"', xy = (212, miny), color = 'g')
        if maxx >= 228.3:
            ax1.axvspan(228.3, 241.5, color = 'g', alpha = 0.1)
        if maxx >= 274.6:
            ax1.axvspan(274.6, 284.6, color = 'g', alpha = 0.1)
            ax1.annotate('Bend A5', xy = (274, miny), color = 'g')
        if maxx >= 309.0:
            ax1.axvspan(309.0, 318.9, color = 'g', alpha = 0.1)
            ax1.annotate('Bend A6', xy = (309, miny), color = 'g')
        if maxx >= 358.8:
            ax1.axvspan(358.8, 370.8, color = 'g', alpha = 0.1)
            ax1.annotate('Bend A7', xy = (359, miny), color = 'g')
        if maxx >= 414.5:
            ax1.axvspan(414.5, 426.5, color = 'g', alpha = 0.1)
            ax1.annotate('Bend A8', xy = (414, miny), color = 'g')
        if maxx >= 445.3:
            ax1.axvspan(445.3, 455.3, color = 'g', alpha = 0.1)
            ax1.annotate('Bend A9', xy = (445, miny), color = 'g')
        if maxx >= 467.5:
            ax1.axvspan(467.5, 479.5, color = 'g', alpha = 0.1)
            ax1.annotate('Bend A10', xy = (467, miny), color = 'g')
        if maxx >= 523.8:
            ax1.axvspan(523.8, 535.8, color = 'g', alpha = 0.1)
            ax1.annotate('Bend A11', xy = (523, miny), color = 'g')
        if maxx >= 571.6:
            ax1.axvspan(571.6, 583.6, color = 'g', alpha = 0.1)
            ax1.annotate('Bend A12', xy = (571, miny), color = 'g')
        if maxx >= 627.9:
            ax1.axvspan(627.9, 639.9, color = 'g', alpha = 0.1)
            ax1.annotate('Bend A13', xy = (627, miny), color = 'g')
        if maxx >= 652.7:
            ax1.axvspan(652.7, 660.7, color = 'g', alpha = 0.1)
            ax1.annotate('Bend A14', xy = (652, miny), color = 'g')
        if maxx >= 678.5:
            ax1.axvspan(678.5, 681.1, color = 'g', alpha = 0.1)
            ax1.annotate('Bend A15', xy = (678, miny), color = 'g')
        
        PLT.title('Average of '+ m + ' IN deployments')
        ax1.grid()
        PLT.errorbar(X_nonull, Y_nonull, xerr = None, yerr = Y_std)
        fig.savefig(outdir + 'average_in_' + m + '.png')
    
    #_____plot all monthly averages together_____#
    fig = PLT.figure(figsize=(15,8),dpi=150)
    box = [0.14, 0.14, 0.76, 0.76]
    ax = fig.add_axes(box)
    ax.set_xlabel('Position (cm)')
    ax.set_ylabel('Torque (in.oz)')
    colorthis = [0, 0, 1]
    denom = int(len(months)) - 3.
    for i, month in enumerate(sorted(months, key = lambda x:datetime.strptime(x,'%b%Y'))):
        if month == 'Oct2012' or month == 'Dec2012' or month == 'Nov2012' or month == 'Jan2013' or month == 'Feb2013' or month == 'Mar2012': continue #  or month == 'Oct2013' because October only has one run
        ax.plot(months[month][0], months[month][1], '-', linewidth = 1, c = tuple(colorthis), label = month)
        colorthis[0] += (1 / denom)
        colorthis[2] -= (1 / denom)
    h, l = ax.get_legend_handles_labels()
    ax.grid()
    PLT.title('Monthly Torque Average (after Feb 2013)')
    PLT.legend(h, l, 'upper left', ncol = 2)
    PLT.savefig(outdir + 'Monthly_averages_Update-' + q + '.png')
    PLT.close(fig)